﻿import random
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from network_env import NetworkEnvironment, draw_network
from matplotlib.animation import FuncAnimation

# 生成随机连通网络  这个程序没用上   
def generate_network(num_nodes=100, connection_prob=0.1):
    while True:
        G = nx.erdos_renyi_graph(num_nodes, connection_prob)
        if nx.is_connected(G):
            return G
    
env = NetworkEnvironment(node_count=25, link_count=76)
G = env.graph
draw_network(G)
nodes = {n: list(G.neighbors(n)) for n in G.nodes()}  # 邻接表
source = env.source
target = env.destination
max_hops = 150 # 根据网络直径自动调整更佳

# 路径生成与验证
def generate_valid_path():
    path = [source]
    current = source
    visited = set()
    
    for _ in range(max_hops):
        if current == target:
            break
        valid_neighbors = [n for n in nodes[current] if n not in visited]
        if not valid_neighbors:
            break
        next_node = random.choice(valid_neighbors)
        path.append(next_node)
        visited.add(next_node)
        current = next_node
    
    # 补齐路径长度
    while len(path) <= max_hops and path[-1] != target:
        path.append(target)
    return path[:max_hops+1]

# 适应度函数加强版（考虑路径有效性）
def enhanced_fitness(path):
    if path[0] != source or path[-1] != target:
        return float('inf')
    
    # 检查路径连通性
    for i in range(len(path)-1):
        if path[i+1] not in nodes[path[i]]:
            return float('inf')
    
    return len(path) - 1  # 跳数优先

# 增强型PSO优化器
class EnhancedPSO:
    def __init__(self, num_particles=100, max_iter=200):
        self.num_particles = num_particles
        self.max_iter = max_iter
        self.gbest = None
        self.fitness_history = []
        
        # 参数自适应调整
        self.w = 0.7
        self.c1 = 1.4
        self.c2 = 1.5
    
    def optimize(self):
        # 粒子初始化
        # swarm是一个列表，包含每个粒子的初始路径、个体最优路径和适应度
        swarm = [{
            'position': generate_valid_path(),
            'pbest': None,  # 个体最优路径
            'pbest_fit': float('inf')
        } for _ in range(self.num_particles)] 
        
        # 初始化全局最优
        self.gbest = min(swarm, 
                       key=lambda x: enhanced_fitness(x['position']))['position'].copy() 
        #这一句理解为先min(swarm,key=lambda x: enhanced_fitness(x['position']))找出符合的字典（swarm的一个元素）
        #再通过['position']取出这个字典的'position'值
        #最后执行.copy()方法，把这个值复制到self.gbest中
        # key是min函数的参数，x是swarm列表中的每个元素

        for iter in range(self.max_iter):
            # 参数衰减
            self.w *= 0.98
            self.c1 *= 0.995
            self.c2 *= 0.995
            
            for particle in swarm:
                current_fit = enhanced_fitness(particle['position'])
                
                # 更新个体最优
                if current_fit < particle['pbest_fit']:
                    particle['pbest'] = particle['position'].copy()
                    particle['pbest_fit'] = current_fit
                
                # 更新全局最优
                if current_fit < enhanced_fitness(self.gbest):
                    self.gbest = particle['position'].copy()
                
                # 速度更新（基于路径相似性）
                new_path = self._update_position(particle)
                particle['position'] = new_path
            
            # 记录收敛过程
            self.fitness_history.append(enhanced_fitness(self.gbest))
            
            print(f"Iteration {iter+1}: Best Hops: {self.fitness_history[-1]}")
        
        return self.gbest
    
    def _update_position(self, particle):
        new_path = [source]
        current_node = source
        
        for step in range(1, max_hops+1):
            if current_node == target:
                break
            
            # 动态候选节点选择
            candidates = []
            neighbors = nodes[current_node]
            
            # 加入各参考方向的节点
            ref_nodes = [
                particle['position'][step] if step < len(particle['position']) else None,
                particle['pbest'][step] if particle['pbest'] and step < len(particle['pbest']) else None,
                self.gbest[step] if step < len(self.gbest) else None
            ]
            
            # 概率选择策略
            for ref in ref_nodes:
                if ref in neighbors:
                    candidates.append(ref)
            
            # 加入随机探索
            if random.random() < 0.2:
                candidates.extend(neighbors)
            
            if not candidates:
                candidates = neighbors.copy()
            
            candidates = list(set(candidates))
            
            # 概率选择（适应度导向）
            if candidates:
                next_node = self._probability_select(candidates, current_node)
                new_path.append(next_node)
                current_node = next_node
            else:
                break
        
        # 路径补齐
        while len(new_path) <= max_hops and new_path[-1] != target:
            new_path.append(target)
        return new_path[:max_hops+1]
    
    def _probability_select(self, candidates, current):
        # 基于到目标节点的距离进行概率加权
        distances = []
        for node in candidates:
            try:
                distance = nx.shortest_path_length(G, node, target)
            except:
                distance = 100
            distances.append(1 / (distance + 1))
        
        total = sum(distances)
        probs = [d/total for d in distances]
        return np.random.choice(candidates, p=probs)

# 可视化模块
class Visualizer:
    @staticmethod # 可视化收敛过程
    def plot_convergence(fitness_history):
        plt.figure(figsize=(10, 6))
        plt.plot(fitness_history, marker='o', linestyle='--')
        plt.title("PSO Convergence Process")
        plt.xlabel("Iteration")
        plt.ylabel("Hops")
        plt.grid(True)
        plt.show()
    
    @staticmethod  # 可视化网络拓扑
    def plot_network(path): 
        plt.figure(figsize=(12, 12))
        pos = nx.spring_layout(G, seed=42)
        
        # 绘制网络拓扑
        nx.draw_networkx_nodes(G, pos, node_color='lightblue')
        nx.draw_networkx_edges(G, pos, edge_color='gray', alpha=0.4)
        
        # 高亮最优路径
        path_edges = [(path[i], path[i+1]) for i in range(len(path)-1)]
        nx.draw_networkx_nodes(G, pos, nodelist=path, node_color='red')
        nx.draw_networkx_edges(G, pos, edgelist=path_edges, 
                              edge_color='red', width=2)
        
        plt.title("Network Topology with Optimized Path")
        plt.show()


def animate_path(G, path):
    fig, ax = plt.subplots(figsize=(12, 12))
    pos = nx.spring_layout(G, seed=42)
    
    # 绘制静态网络
    nx.draw_networkx_nodes(G, pos, node_color='lightblue', ax=ax)
    nx.draw_networkx_edges(G, pos, edge_color='gray', alpha=0.4, ax=ax)
    
    # 高亮最优路径（静态部分）
    path_edges = list(zip(path, path[1:]))
    nx.draw_networkx_nodes(G, pos, nodelist=path, node_color='red', ax=ax)
    nx.draw_networkx_edges(G, pos, edgelist=path_edges, 
                          edge_color='red', width=2, ax=ax)
    
    # 创建动态红点（初始位置空列表）
    red_dot, = ax.plot([], [], 'o', color='blue', 
                      markersize=15, animated=True)
    
    # 预提取路径节点坐标（转换为Nx2数组）
    node_coords = np.array([pos[node] for node in path])  # 形状：(N,2)

    def animate(i):
        # 关键修复：将标量包装成单元素列表
        x, y = node_coords[i]  # 获取当前坐标
        red_dot.set_data([x], [y])  # 包装成列表传递
        return red_dot,

    ani = FuncAnimation(
        fig, animate,
        frames=len(path),
        interval=500,
        blit=True,
        repeat=True
    )
    plt.title("Dynamic Path Visualization")
    plt.show()
    return ani


# 主程序
if __name__ == "__main__":
    # 优化过程
    pso = EnhancedPSO(num_particles=150, max_iter=40)
    best_path = pso.optimize()
    
    # 结果可视化
    print("\nOptimized Path:", best_path)
    print(f"Total Hops: {len(best_path)-1}")
    
    #Visualizer.plot_convergence(pso.fitness_history)
    #Visualizer.plot_network(best_path)
    ani = animate_path(G, best_path)
